import os.path

import paddlex as pdx
from paddlex import transforms as T
from src.config.config import DATA_ROOT, OUTPUT_DIR

train_transforms = T.Compose([
    T.MixupImage(mixup_epoch=250), T.RandomDistort(),
    T.RandomExpand(im_padding_value=[123.675, 116.28, 103.53]), T.RandomCrop(),
    T.RandomHorizontalFlip(), T.BatchRandomResize(
        target_sizes=[320, 352, 384, 416, 448, 480, 512, 544, 576, 608],
        interp='RANDOM'), T.Normalize(
        mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

eval_transforms = T.Compose([
    T.Resize(
        608, interp='CUBIC'), T.Normalize(
        mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

base_path = os.path.join(DATA_ROOT, 'insect_det')
train_dataset = pdx.datasets.VOCDetection(
    data_dir=base_path,
    file_list=os.path.join(base_path, 'train_list.txt'),
    label_list=os.path.join(base_path, 'labels.txt'),
    transforms=train_transforms,
    shuffle=True)
eval_dataset = pdx.datasets.VOCDetection(
    data_dir=base_path,
    file_list=os.path.join(base_path, 'val_list.txt'),
    label_list=os.path.join(base_path, 'labels.txt'),
    transforms=eval_transforms)

num_classes = len(train_dataset.labels)
model = pdx.det.YOLOv3(num_classes=num_classes, backbone='DarkNet53', label_smooth=True)  # 提升小样本泛化能力
model.train(
    num_epochs=270,
    train_dataset=train_dataset,
    train_batch_size=1,
    eval_dataset=eval_dataset,
    learning_rate=0.000125,
    lr_decay_epochs=[216, 243],
    warmup_steps=250,
    warmup_start_lr=0.0,
    save_interval_epochs=5,
    save_dir=os.path.join(OUTPUT_DIR, 'yolov3_darknet53'),
    use_vdl=True)
